#include <iostream>
#include <string>
#include <qpOASES.hpp>
#include "MPCsolver.h"

using namespace std;
using namespace Eigen; // conflict with namespace "qpOASES"
using qpOASES::QProblem;
using qpOASES::real_t;

MPCsolver::MPCsolver()
{ 
    // A_qp_ = Matrix<double, 13 * PREDICTION_HORIZON, 13>();
    // B_qp_ = Matrix<double, 13 * PREDICTION_HORIZON, 12 * PREDICTION_HORIZON>();
    // Q_ = Matrix<double, 13 * PREDICTION_HORIZON, 13 * PREDICTION_HORIZON>();
    // R_ = Matrix<double, 12 * PREDICTION_HORIZON, 12 * PREDICTION_HORIZON>();
    // D_ = Vector<double, 13 * PREDICTION_HORIZON>();

    // H_ = Matrix<double, 12 * PREDICTION_HORIZON, 12 * PREDICTION_HORIZON>();
    // g_ = Vector<double, 12 * PREDICTION_HORIZON>();

    // U_ = Vector<double, 12 * PREDICTION_HORIZON>();

    Vector<double, 13> Q_tmpvec;
    Q_tmpvec << 25, 25, 10, 1, 1, 100, 0, 0, 0.3, 0.2, 0.2, 20, 0;
    Q_ = Q_tmpvec.replicate(PREDICTION_HORIZON, 1).asDiagonal();  // use diagonalMatrix to accelerate calculation

    Vector<double, 12> R_tmpvec;
    R_tmpvec << 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005, 0.00005;
    R_ = R_tmpvec.replicate(PREDICTION_HORIZON, 1).asDiagonal(); // use diagonalMatrix to accelerate calculation

    c_lb_.setZero(); // 0(20h*1)
    
    Vector<double, 5> c_i_ub;
    c_i_ub << INF, INF, INF, INF, f_max;
    c_ub_ = c_i_ub.replicate(PREDICTION_HORIZON, 1); // [INF, INF, INF, INF, f_max](20h*1)

}

MPCsolver::~MPCsolver()
{
}

void MPCsolver::calA_qp_and_B_qp(const Vector<double, PREDICTION_HORIZON> &psi_k_d,
                                 const Matrix<double, 3 * PREDICTION_HORIZON, 3> &r_k_ix)
{
    // I_B数组映射为Eigen矩阵
    Map<Matrix3d> I_B(I_B_array[0]);

    // TODO:r_k_ix分解为1~4
    Matrix3d r_k_1x;
    Matrix3d r_k_2x;
    Matrix3d r_k_3x;
    Matrix3d r_k_4x;

    int h = PREDICTION_HORIZON;
    int delt = DELTA_T_MPC;
    Matrix<double, 13, 13> A_k;
    Matrix3d tmpmat;
    // clang-format off
    A_k <<  1, 0, 0, 0, 0, 0, -1, -1, -1, 0, 0, 0, 0,
            0, 1, 0, 0, 0, 0, -1, -1, -1, 0, 0, 0, 0,
            0, 0, 1, 0, 0, 0, -1, -1, -1, 0, 0, 0, 0,
            0, 0, 0, 1, 0, 0, 0, 0, 0, 0, delt, 0, 0, 0,
            0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, delt, 0, 0,
            0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, delt, 0,
            0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, delt,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0;
    // clang-format on

    // 计算A(0)并赋值给A_qp
    A_k.block<3, 3>(0, 6) = (R("z", psi_k_d(0))).transpose() * delt;
    A_qp_.block<13, 13>(0, 0) = A_k;

    Matrix<double, 13, 12> B_k;
    B_k.setZero();

    Matrix3d R_k_B_O = R("z", psi_k_d(0));
    Matrix3d I_k_p_inverse = ((R_k_B_O * I_B * R_k_B_O.transpose()).inverse());

    //  计算B(0)并赋值给B_qp
    B_k.block<3, 3>(6, 0) = I_k_p_inverse * r_k_1x * delt;
    B_k.block<3, 3>(6, 3) = I_k_p_inverse * r_k_2x * delt;
    B_k.block<3, 3>(6, 6) = I_k_p_inverse * r_k_3x * delt;
    B_k.block<3, 3>(6, 9) = I_k_p_inverse * r_k_4x * delt;
    Matrix3d I3deltm = Matrix3d::Identity() * delt / Mass;
    for (int _i = 0; _i < 4; _i++)
        B_k.block<3, 3>(9, 3 * _i) = I3deltm;

    B_qp_.block<13, 12>(0, 0) = B_k;

    for (int i = 1; i < h; i++)
    {
        R_k_B_O = R("z", psi_k_d(i));
        I_k_p_inverse = ((R_k_B_O * I_B * R_k_B_O.transpose()).inverse());

        // 计算A(1)~A(h-1)并赋值给A_qp
        A_k.block<3, 3>(0, 6) = R_k_B_O.transpose() * delt;
        A_qp_.block<13, 13>(i * 13, 0) = A_k * A_qp_.block<13, 13>((i - 1) * 13, 0);

        // 计算B(1)~B(h-1)并赋值给B_qp

        B_k.block<3, 3>(6, 0) = I_k_p_inverse * r_k_1x * delt;
        B_k.block<3, 3>(6, 3) = I_k_p_inverse * r_k_2x * delt;
        B_k.block<3, 3>(6, 6) = I_k_p_inverse * r_k_3x * delt;
        B_k.block<3, 3>(6, 9) = I_k_p_inverse * r_k_4x * delt;

        Matrix3d I3deltm = Matrix3d::Identity() * delt / Mass;
        for (int _i = 0; _i < 4; _i++)
            B_k.block<3, 3>(9, 3 * _i) = I3deltm;

        B_qp_.block<13, 12>(13 * i, 12 * i) = B_k;

        for (int j = 0; j < i; j++)
        {
            B_qp_.block<13, 12>(13 * i, 12 * j) =
                A_k * B_qp_.block<13, 12>(13 * (i - 1), 12 * j);
        }
    }
}

void MPCsolver::setD_and_x0(const Vector<double, 13 * PREDICTION_HORIZON> &vecD, const Vector<double, 13> &x_0)
{
    D_ = vecD;
    x_0_ = x_0;
}

void MPCsolver::calH_and_g()
{
    H_ = 2 * ((B_qp_.transpose() * Q_ * B_qp_) + R_);
    g_ = 2 * B_qp_.transpose() * Q_ * (A_qp_ * x_0_ - D_);
}

void MPCsolver::qpSolveU()
{
    
    int h = PREDICTION_HORIZON;
    QProblem qproblem(12 * h, 20 * h); // 优化变量维度nV, 约束变量维度nC

    real_t* H_qp = H_.data();
    real_t* g_qp = g_.data(); // 线性项系数 与maatlab不同为 x' * g
    real_t* C_qp = C_.data();    // 不等式/等式约束系数矩阵
    real_t* lb_qp;
    real_t* ub_qp;
    real_t* lbC_qp = c_lb_.data(); // 不等式/等式上界 上下界相等时为等式约束
    real_t* ubC_qp = c_ub_.data(); // 不等式/等式下界

    real_t U_opt[12 * h];
    int nWSR = 10; // 最大迭代次数

    cout<<qproblem.init(H_qp, g_qp, C_qp, lb_qp, ub_qp, lbC_qp, ubC_qp, nWSR)<<endl;
    qproblem.getPrimalSolution(U_opt);

}
